import os
import sys
import wandb

# 获取当前脚本路径信息（解决模块导入和配置文件路径问题）
script_path = os.path.abspath(__file__)
script_dir = os.path.dirname(script_path)
project_root = os.path.dirname(script_dir)

# 将项目根目录添加到Python搜索路径
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# 导入项目模块
from training.utils.config_util import load_config
from training.kfold_trainer import main as kfold_trainer
from training.utils.common_util import set_random_seed


def merge_nested_dicts(base, update):
    """
    递归合并字典（支持嵌套结构）：update参数覆盖base同名参数，不破坏嵌套层级
    
    参数:
        base (dict): 基础字典
        update (dict): 更新字典
    
    返回:
        dict: 合并后的字典
    """
    merged = base.copy()
    
    for key, value in update.items():
        if isinstance(value, dict) and key in merged and isinstance(merged[key], dict):
            # 递归合并嵌套字典
            merged[key] = merge_nested_dicts(merged[key], value)
        else:
            # 直接覆盖或添加值
            merged[key] = value
    
    return merged


def validate_optimizer_config(config):
    """
    验证并清理优化器配置
    
    
    参数:
        config (dict): 完整配置字典
    
    返回:
        dict: 验证后的配置字典
    """
    if 'optimizer' in config and 'args' in config['optimizer']:
        optimizer_type = config['optimizer'].get('type', '')
        optimizer_args = config['optimizer']['args']
        
        # 记录当前的优化器配置
        print(f"当前优化器类型: {optimizer_type}")
        print(f"当前优化器参数: {optimizer_args}")
        
        # 注意：详细的参数验证和过滤现在由optimizer_util模块处理
        # 这里只做基本的日志记录
    
    return config


def main():
    """
    主函数：加载配置→合并超参→验证配置→启动训练
    单进程模式运行，不使用DDP分布式训练
    """
    print("开始执行超参数调优，单进程模式")
    
    try:
        # 1. 加载本地固定配置
        config_path = os.path.join(script_dir, "config.yaml")
        config = load_config(config_path)
        
        # 确保trainer部分存在
        if 'trainer' not in config:
            config['trainer'] = {}
        
        # 设置随机种子，确保实验可复现
        seed = config.get('random_seed', 42)
        set_random_seed(seed)
        
        # 2. 初始化wandb，获取sweep动态生成的超参数
        wandb.init(project=config['wandb']['project'],
                   dir=project_root,
                   resume=True,
                   config=config,
                   mode=config['wandb']['mode']
                   )
        print(f"当前wandb运行ID: {wandb.run.id}")
        
        # 设置保存目录
        save_dir = os.path.join(script_dir, wandb.run.id)
        config['trainer']['save_dir'] = save_dir
        os.makedirs(save_dir, exist_ok=True)
        
        # 获取sweep参数
        sweep_params = wandb.config
        print(f"从wandb获取的参数: {sweep_params}")
        
        # 合并参数
        final_config = merge_nested_dicts(config, sweep_params)
        
        # 验证优化器配置
        final_config = validate_optimizer_config(final_config)
        
        # 记录最终配置
        print("\n✅ 最终合并配置:")
        for section, values in final_config.items():
            print(f"{section}: {values}")
        
        # 启动k折训练
        print("开始k折训练...")
        kfold_trainer(final_config)
        
    except Exception as e:
        # 记录错误
        print(f"[错误] 训练过程中出错: {str(e)}")
        raise


if __name__ == "__main__":
    # 单次运行仅执行1次主函数
    main()